package cn.classba.hive.io.format;

import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.exec.FileSinkOperator;
import org.apache.hadoop.hive.ql.exec.Utilities;
import org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat;
import org.apache.hadoop.hive.serde.serdeConstants;
import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.util.Progressable;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Properties;

public class MultiDelimitHiveIgnoreKeyTextOutputFormat extends HiveIgnoreKeyTextOutputFormat {

    @Override
    public FileSinkOperator.RecordWriter getHiveRecordWriter(JobConf jc, Path outPath, Class valueClass, boolean isCompressed, Properties tableProperties, Progressable progress) throws IOException {
        String rowSeparatorString = tableProperties.getProperty(
                serdeConstants.LINE_DELIM, "\n");

        final byte[] recordDelimiterBytes = rowSeparatorString.getBytes("UTF-8");
        FileSystem fs = outPath.getFileSystem(jc);
        final OutputStream outStream = Utilities.createCompressedStream(jc,
                fs.create(outPath, progress), isCompressed);
        return new FileSinkOperator.RecordWriter() {
            @Override
            public void write(Writable r) throws IOException {
                if (r instanceof Text) {
                    Text tr = (Text) r;
                    outStream.write(tr.getBytes(), 0, tr.getLength());
                    outStream.write(recordDelimiterBytes);
                } else {
                    // DynamicSerDe always writes out BytesWritable
                    BytesWritable bw = (BytesWritable) r;
                    outStream.write(bw.get(), 0, bw.getSize());
                    outStream.write(recordDelimiterBytes);
                }
            }

            @Override
            public void close(boolean abort) throws IOException {
                outStream.close();
            }
        };
    }
}
